import logging
from typing import List, TypeVar, Optional

from typing import Any, Dict
from calc_manager.algorithm.algo import Algorithm
from calc_manager.base.executor import TaskExecutor
from calc_manager.base.runnable import Runnable
from calc_manager.base.task import (DefaultRunnableTask, Task, TaskFactory,
                                    TaskInfo)
from calc_manager.operator import Operator, OperatorInfo, print_operator_info
from calc_manager.plugins.plugins_manager import PluginManager  

logger = logging.getLogger(__name__)


class DefaultExecutor(TaskExecutor):
    def initialize(self):
        pass

    def shutdown(self):
        pass

    def invoke(self, task: Task):
        return task.execute()

    def invoke_tasks(self, tasks: List[Task]):
        results = {}
        current_ctx: Dict[str, Any] = {}
        for task in tasks:
            try:
                if current_ctx:
                    task.set_ctx(current_ctx)
                ret = self.invoke(task)
                current_ctx = task.get_ctx()
                
                # If ret is a dict, merge it into results
                if isinstance(ret, dict):
                    results.update(ret)
                else:
                    # Use task name as key if ret is not a dict
                    results[task.task_info.task_name] = ret
                    
            except Exception as err:
                logger.error("Failed to run task {%s}: {%s}", task.task_info.task_name, err)
                break

        return results

class AlgorithmRunnable(Runnable):
    def __init__(self, algo: Algorithm):
        super().__init__(algo.info.algo_name)
        self.algo = algo

    def set_operator_runnable_name(self, name):
        self.set_runnable_name(name)

    def _run(self, *argc, **argv):
        try:
            return self.algo.algo_run_task()
        except Exception as err:
            logger.error(f"Failed to run algorighm {self.algo.info}: {err}")
            raise err

class OperatorRunnable(Runnable):
    def __init__(self, operator_info: OperatorInfo, *argc, **argv):
        super().__init__(operator_info.operator_name, *argc, **argv)
        self.operator_info = operator_info

    def set_operator_runnable_name(self, name):
        self.set_runnable_name(name)

    def _run(self, *argc, **argv):
        operator = Operator(self.operator_info)
        #logger.debug("oerator has fully initializea..")
        #logger.debug({print_operator_info(operator)})
        try:
            # mix argv and extra_argv
            argv = {**argv, **self.extra_argv}
            logger.debug("Running operator %s, argc count: %d, argv keys: %s" % (self.operator_info.operator_name, len(argc), list(argv.keys())))
            # extend argv to extra_argv
            return operator.invoke_func(*argc, **argv)
        except Exception as err:
            logger.error(f"Failed to run operator {self.operator_info.operator_name}: {err}")
            raise err
        finally:
            operator.dispose()


def _get_runnable_instance_from_operator(
    oi: OperatorInfo, *argc, **argv
) -> OperatorRunnable:
    runnable_operator = OperatorRunnable(oi)
    runnable_operator.set_run_param(*argc, **argv)
    return runnable_operator


def operator_exec_runnable(
    operator_info: OperatorInfo, *argc, **argv
) -> OperatorRunnable:
    return _get_runnable_instance_from_operator(operator_info, *argc, **argv)


def exec_operator_task(
    executor: TaskExecutor,
    operator_info: OperatorInfo,
    task_factory: TaskFactory[Task],
    plugins_argv: Dict[str, Any] = {}, # plugin_name + plugin_param
    plugin_manager: Optional[PluginManager] = None,
    *argc,
    **argv,
):
    task = build_operator_task_with_args(
        "test-task-name", operator_info, task_factory, plugins_argv, plugin_manager, *argc, **argv
    )
    return executor.invoke(task)


#def get_plugin_invoke_param(plugin_manager: PluginManager, plugins_argv: Dict[str, Any] = {}):
#    invoke_param = {}
#    if plugin_manager:
#        invoke_param = Operator.get_plugin_objects_from_argv(plugin_manager, plugins_argv)
#        invoke_param = {k: v.get_plugin_ojbect() for k, v in invoke_param.items()}
#    return invoke_param

def build_operator_task_with_args(
    task_name: str,
    operator_info: OperatorInfo,
    task_factory: TaskFactory[Task],
    plugins_argv: Dict[str, Any] = {},
    plugin_manager: Optional[PluginManager] = None,
    *argc,
    **argv,
) -> Task:
    op_runnable = operator_exec_runnable(operator_info, *argc, **argv)

    #plugin invoke
    #invoke_param = get_plugin_invoke_param(plugin_manager, plugins_argv)
    #op_runnable.set_extra_param(**invoke_param)

    return task_factory.create_task(TaskInfo.create_task_info(op_runnable, task_name))


def build_operator_task_without_args(
    task_name: str,
    operator_info: OperatorInfo,
    task_factory: TaskFactory[Task],
) -> Task:
    return build_operator_task_with_args(task_name, operator_info, task_factory)

def build_algo_single_task(task_factory: TaskFactory[Task], algo: Algorithm) -> Task:
    """Build a single task for the entire algorithm execution
    
    Args:
        task_factory: Factory to create tasks
        algo: Algorithm to execute
        
    Returns:
        Task: Single task that will execute the entire algorithm
    """
    algo_runnable = AlgorithmRunnable(algo)
    return task_factory.create_task(TaskInfo.create_task_info(algo_runnable, algo.info.algo_name))
    

def build_algo_tasks(task_factory: TaskFactory[Task], algo: Algorithm) -> List[Task]:
    tasks = []
    for op_graph_id, opinfo in algo.get_operation_info_exection_order_by_graphid():
        # not set param here because param will set by algorithm 
        # plugins will be decoded and handled by algorithm
        op_runnable = operator_exec_runnable(opinfo)
        op_runnable.set_operator_runnable_name(op_graph_id)
        task_name = f"task-{op_graph_id}, {opinfo.operator_name}"
        task_info = TaskInfo.create_task_info(op_runnable, task_name)
        logger.debug("create task by info %s", task_info)
        task = task_factory.create_task(task_info)
        tasks.append(task)
    return tasks


#TODO: fix ref by remote_task
#TODO: fix ref by remote_task
#TODO: fix ref by remote_task
class AlgoTaskFactory(TaskFactory[DefaultRunnableTask]):
    def __init__(self, algo: Algorithm):
        self.algo = algo

    def create_task(self, info: TaskInfo) -> DefaultRunnableTask:
        task = DefaultRunnableTask(info)

        def pre_run(task: DefaultRunnableTask):
            ctx = task.get_ctx()
            logger.debug("pre_run task ctx: %s", ctx)
            argv = self.algo.algo_params
            self.algo.mapping_all_plugins_to_extra_param(
                ctx, task.runnable, argv
            )

            #pop out raw plugin info
            argv.pop(Operator.INVOKE_PLUGINS_KEY, None)

            self.algo.mapping_param_to_ctx_argv(
                ctx, task.runnable, argv
            )
            logger.debug("pre_run task finished, argv : %s", argv)

        def post_run(task: DefaultRunnableTask, ret):
            ctx = task.get_ctx()
            self.algo.mapping_result_to_ctx(ctx, task.runnable, ret)
            ctx = task.get_ctx()
            logger.debug("post_run task ctx: %s", ctx)

        task.add_pre_run_callable(pre_run)
        task.add_post_run_callable(post_run)

        return task
